# coding=utf-8
import os

from django.forms import ValidationError, ModelForm

from lib import sina
from .models import BackupStock, LiveStock

codes = []
if os.path.exists("./codes"):
    codes = open("./codes").read().splitlines()


def clean_code(code: str):
    if len(code) != 6 and len(code) != 8:
        raise ValidationError("代码格式错误")

    if code.startswith("0") or code.startswith("3"):
        code = "sz" + code
    elif code.startswith("6"):
        code = "sh" + code
    elif code.startswith("sz") or code.startswith("sh"):
        code = code
    else:
        raise ValidationError("不支持的代码")

    if code in codes:
        return code
    else:
        try:
            sina.get_stock_data([code])
        except Exception as e:
            raise ValidationError(f"代码验证失败：{e}")

        # 将验证过的代码缓存到磁盘，每次启动加载到内存
        codes.append(code)
        open("./codes", "a+").write(f"{code}\n")
        return code


class BackupStockForm(ModelForm):
    class Meta:
        model = BackupStock
        exclude = []

    def clean_code(self):
        code = self.cleaned_data.get("code")

        return clean_code(code)


class LiveStockForm(ModelForm):
    class Meta:
        model = LiveStock
        exclude = []

    def clean_code(self):
        code = self.cleaned_data.get("code")

        return clean_code(code)
